#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <openssl/evp.h>
#include <openssl/pem.h>

// 复用 sm2_encrypt 的核心逻辑（仅用于演示）
int sm2_encrypt_internal(EVP_PKEY *pubkey, const unsigned char *in, size_t in_len,
                         unsigned char **out, size_t *out_len) {
    EVP_PKEY_CTX *ctx = EVP_PKEY_CTX_new(pubkey, NULL);
    if (!ctx) return -1;
    if (EVP_PKEY_encrypt_init(ctx) <= 0 ||
        EVP_PKEY_CTX_set_ec_scheme(ctx, NID_sm2) <= 0 ||
        EVP_PKEY_encrypt(ctx, NULL, out_len, in, in_len) <= 0) {
        EVP_PKEY_CTX_free(ctx); return -1;
    }
    *out = malloc(*out_len);
    int ret = EVP_PKEY_encrypt(ctx, *out, out_len, in, in_len);
    EVP_PKEY_CTX_free(ctx);
    return ret <= 0 ? -1 : 0;
}

int sm2_decrypt(const char *privkey_file, const unsigned char *in, size_t in_len,
                unsigned char **out, size_t *out_len) {
    FILE *fp = fopen(privkey_file, "r");
    EVP_PKEY *pkey = PEM_read_PrivateKey(fp, NULL, NULL, NULL);
    fclose(fp);
    if (!pkey) return -1;

    EVP_PKEY_CTX *ctx = EVP_PKEY_CTX_new(pkey, NULL);
    if (!ctx) { EVP_PKEY_free(pkey); return -1; }

    if (EVP_PKEY_decrypt_init(ctx) <= 0 ||
        EVP_PKEY_CTX_set_ec_scheme(ctx, NID_sm2) <= 0 ||
        EVP_PKEY_decrypt(ctx, NULL, out_len, in, in_len) <= 0) {
        EVP_PKEY_CTX_free(ctx); EVP_PKEY_free(pkey); return -1;
    }

    *out = malloc(*out_len);
    if (EVP_PKEY_decrypt(ctx, *out, out_len, in, in_len) <= 0) {
        free(*out); *out = NULL;
        EVP_PKEY_CTX_free(ctx); EVP_PKEY_free(pkey); return -1;
    }

    EVP_PKEY_CTX_free(ctx); EVP_PKEY_free(pkey);
    return 0;
}

int main() {
    const char *msg = "fsl";

    // 读取公钥用于加密
    FILE *fp = fopen("sm2public_key.pem", "r");
    EVP_PKEY *pubkey = PEM_read_PUBKEY(fp, NULL, NULL, NULL);
    fclose(fp);
    if (!pubkey) { fprintf(stderr, "Load public key failed\n"); return 1; }

    // 加密
    unsigned char *ciphertext = NULL;
    size_t cipher_len = 0;
    if (sm2_encrypt_internal(pubkey, (unsigned char*)msg, strlen(msg), &ciphertext, &cipher_len) != 0) {
        EVP_PKEY_free(pubkey);
        fprintf(stderr, "Encrypt failed\n");
        return 1;
    }
    EVP_PKEY_free(pubkey);

    // 解密
    unsigned char *plaintext = NULL;
    size_t plain_len = 0;
    if (sm2_decrypt("sm2private_key.pem", ciphertext, cipher_len, &plaintext, &plain_len) == 0) {
        plaintext[plain_len] = '\0';
        printf("SM2 decrypt result: %s\n", plaintext);
        free(plaintext);
    } else {
        fprintf(stderr, "SM2 decrypt failed\n");
        free(ciphertext);
        return 1;
    }

    free(ciphertext);
    return 0;
}
